package cn.st4rlight.util.excel;

import java.io.ByteArrayOutputStream;
import java.util.List;
import java.util.Map;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.time.StopWatch;

import com.alibaba.excel.EasyExcelFactory;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.builder.ExcelWriterSheetBuilder;
import com.alibaba.excel.write.handler.WriteHandler;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.google.common.collect.Lists;

import cn.st4rlight.util.collection.TransformUtil;
import cn.st4rlight.util.excel.context.DynamicSheetContext;
import cn.st4rlight.util.excel.strategy.AutoWidthStrategy;
import cn.st4rlight.util.exception.ServiceException;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * @author st4rlight <st4rlight@163.com>
 * Created on 2024-04-14
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class DynamicExcelUtil {

    // 单个sheet最多100w行（excel允许的上限是104_8576）
    public static final int MAX_ROWS_PER_SHEET = 80_0000;


    /**
     * 构建动态列，单个sheet的excel文件
     * NOTE: 单个sheet数据不超过100w行
     *
     * @param context 上下文参数信息
     * @return excel文件字节数组
     */
    public static byte[] buildSingleDynamicExcel(DynamicSheetContext context) {
        return DynamicExcelUtil.buildMultiDynamicExcel(Lists.newArrayList(context));
    }

    /**
     * 构建动态列，多个sheet的excel文件
     * NOTE: 单个sheet数据不超过100w行
     *
     * @param contextList 上下文参数信息列表
     * @return excel文件字节数组
     */
    public static byte[] buildMultiDynamicExcel(List<DynamicSheetContext> contextList) {
        return DynamicExcelUtil.buildMultiDynamicExcel(contextList, Lists.newArrayList());
    }

    /**
     * 构建动态列，多个sheet的excel文件
     * NOTE: 单个sheet数据不超过100w行
     *
     * @param contextList 上下文参数信息列表
     * @param handlerList 定制化到处策略
     * @return excel文件字节数组
     */
    public static byte[] buildMultiDynamicExcel(List<DynamicSheetContext> contextList, List<WriteHandler> handlerList) {
        if (CollectionUtils.isEmpty(contextList)) {
            throw ServiceException.of("用于构建excel的上下文参数列表不能为空");
        }

        // 拆分超过100w的context
        List<DynamicSheetContext> newContextList = DynamicExcelUtil.splitLargeContext(contextList);
        log.info("buildMultiDynamicExcel sheet数量: {}", newContextList.size());

        // 开始处理excel导出
        StopWatch stopWatch = StopWatch.createStarted();
        try(
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ExcelWriter excelWriter = EasyExcelFactory.write(bos).build()
        ) {
            // 逐个转成sheet
            for (int i = 0; i < newContextList.size(); i++) {
                log.info("buildMultiDynamicExcel {} / {}", i + 1, newContextList.size());
                DynamicSheetContext param = newContextList.get(i);
                if (CollectionUtils.isEmpty(param.getHeadKeyList())) {
                    throw ServiceException.of("用于构建excel的表头信息不能为空");
                }

                // 将数据转成 List<List<String>> 格式
                List<List<String>> dataList = DynamicExcelUtil.transformDynamicData(param);
                WriteSheet writeSheet = DynamicExcelUtil.writeSingleSheet(i, param, handlerList);
                // 写出当前sheet到输出流
                excelWriter.write(dataList, writeSheet);
            }

            // 内部会关闭流
            excelWriter.finish();
            log.info("buildMultiDynamicExcel完成 耗时: {}ms, size: {}", stopWatch.getTime(), bos.size());
            return bos.toByteArray();
        } catch (Exception ex) {
            log.error("动态导出excel文件异常", ex);
            throw ServiceException.of(ex.getMessage());
        }
    }

    /**
     * 单个sheet最多100w行，拆分context
     */
    public static List<DynamicSheetContext> splitLargeContext(List<DynamicSheetContext> contextList) {
        List<DynamicSheetContext> newContextList = Lists.newArrayList();

        for (DynamicSheetContext context : contextList) {
            List<Map<String, String>> rowList = context.getRowList();
            if (rowList.size() <= MAX_ROWS_PER_SHEET) {
                newContextList.add(context);
                continue;
            }

            // 超过100w行拆分成多个sheet
            List<List<Map<String, String>>> partitionList = Lists.partition(rowList, MAX_ROWS_PER_SHEET);
            for (List<Map<String, String>> partition : partitionList) {
                DynamicSheetContext newContext = new DynamicSheetContext()
                        .setHeadKeyList(context.getHeadKeyList())
                        .setHeadNameList(context.getHeadNameList())
                        .setRowList(partition)
                        .setSheetName(context.getSheetName());
                newContextList.add(newContext);
            }
        }

        log.info("splitLargeContext origin size: {}, result size: {}", contextList.size(), newContextList.size());
        return newContextList;
    }

    /**
     * 写出参数到单个sheet
     *
     * @param sheetIndex sheet在excel文件中的下标
     * @param context 上下文参数
     * @param handlerList 定制化写出策略
     * @return 单个sheet
     */
    private static WriteSheet writeSingleSheet(
        int sheetIndex, DynamicSheetContext context, List<WriteHandler> handlerList
    ) {
        // header需要 List<List<String>> 形式, 比如[[a], [b], [c]]
        List<String>   headNameList = context.getHeadNameList();
        List<List<String>> headList = TransformUtil.toList(headNameList, Lists::newArrayList);

        ExcelWriterSheetBuilder builder = EasyExcelFactory.writerSheet(sheetIndex, context.getSheetName())
                .registerWriteHandler(new AutoWidthStrategy())
                .head(headList);

        if (CollectionUtils.isNotEmpty(handlerList)) {
            handlerList.forEach(builder::registerWriteHandler);
        }

        return builder.build();
    }

    /**
     * 将数据转成 List<List<String>> 格式
     */
    private static List<List<String>> transformDynamicData(DynamicSheetContext context) {
        List<String> headKeyList = context.getHeadKeyList();
        List<Map<String, String>> rowList = context.getRowList();

        // 转成 List<List<String>> 格式
        List<List<String>> dataList = Lists.newArrayList();
        for (Map<String, String> singleRow : rowList) {
            List<String> itemList = TransformUtil.toList(headKeyList, headKey -> singleRow.getOrDefault(headKey, ""));
            dataList.add(itemList);
        }

        return dataList;
    }
}
